import os

import torch
import cv2
import numpy as np
from models.experimental import attempt_load  # 加载 YOLOv7 的架构
from utils.general import non_max_suppression, scale_coords

# 1. 加载 YOLOv7 模型
def load_model(weights_path='best.pt', device='cuda'):
    # 加载模型架构及权重
    model = attempt_load(weights_path, map_location=device)
    model.eval()  # 设置为评估模式
    return model

# 2. 加载并预处理图片
def load_image(image_path, img_size=640):
    image = cv2.imread(image_path)  # 使用 OpenCV 读取图片
    original_image = image.copy()  # 保留原图用于后续绘制结果
    image = cv2.resize(image, (img_size, img_size))  # 将图片调整到 640x640
    image = image[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB and HWC to CHW
    image = np.ascontiguousarray(image)  # 使数组在内存中连续
    image = torch.from_numpy(image).float() / 255.0  # 转换为张量并归一化
    image = image.unsqueeze(0)  # 增加批次维度
    return image, original_image

# 3. 对图片进行推理
def infer_image(model, image, device='cuda'):
    image = image.to(device)  # 将图片移动到 GPU 或 CPU
    with torch.no_grad():  # 禁用梯度计算，以提高推理速度
        outputs = model(image)[0]  # 使用模型进行推理，获取第一个输出（检测框）
    return outputs

# 4. 进行非极大值抑制 (NMS) 处理
def process_detections(outputs, conf_threshold=0.25, iou_threshold=0.45):
    # 使用非极大值抑制来消除重叠框
    results = non_max_suppression(outputs, conf_threshold, iou_threshold)
    return results

# 5. 在图片上绘制检测结果
def draw_detections(original_image, results, img_size=640):\
    h, w = original_image.shape[:2]
    for result in results:
        if result is not None:
            result[:, :4] = scale_coords(image.shape[2:], result[:, :4], original_image.shape).round()
            for detection in result:
                x1, y1, x2, y2, conf, cls = detection[:6].cpu().numpy()
                # 绘制检测框
                cv2.rectangle(original_image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
                # 显示类别和置信度
                cv2.putText(original_image, f'Class: {int(cls)}, Conf: {conf:.2f}',
                            (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
    return original_image

def yolo_detect(image_folder,output_path):
        weights_path = 'best.pt'  # 替换为你的模型权重路径
        device = 'cuda' if torch.cuda.is_available() else 'cpu'

        # 加载模型
        model = load_model(weights_path, device)
        for image in os.listdir(image_folder):
            image_path = os.path.join(image_folder, image)
            output_image_path = os.path.join(output_path, image)
            img_size = 640
            image, original_image = load_image(image_path, img_size)
            # 推理
            outputs = infer_image(model, image, device)

            # 后处理：非极大值抑制
            results = process_detections(outputs)

            # 在图片上绘制检测框
            draw_detections(original_image, results)
            cv2.imwrite(output_image_path, result_image)

